# @file    : multimodal_content_to_pdf
# @time    : 2025/4/14
# @author  : yongpeng.yao
# @desc    :
import os

from PIL import Image as PILImage
from dotenv import load_dotenv
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate
from langchain_openai import ChatOpenAI
from reportlab.lib import colors
from reportlab.lib.styles import ParagraphStyle
from reportlab.lib.styles import getSampleStyleSheet
from reportlab.lib.units import inch
from reportlab.pdfbase import pdfmetrics
from reportlab.pdfbase.ttfonts import TTFont
from reportlab.platypus import Image, PageBreak
from reportlab.platypus import SimpleDocTemplate, Paragraph, Spacer

from generate_pdf.image_analyse import image_analyse_to_text

# 加载环境变量
load_dotenv()


class ContentOptimizer:
    def __init__(self):
        # 初始化文本优化模型
        self.text_llm = ChatOpenAI(model="qwen-max", api_key=os.getenv("DASHSCOPE_API_KEY"),
                                   base_url="https://dashscope.aliyuncs.com/compatible-mode/v1")

        # 设置输出解析器
        self.output_parser = StrOutputParser()

        # 定义样式
        self.styles = self._init_styles()

    @staticmethod
    def _init_styles():
        """初始化PDF样式"""
        font_name = setup_chinese_fonts()
        styles = getSampleStyleSheet()
        # 添加自定义样式
        styles.add(ParagraphStyle(
            name="Content",
            fontName=font_name,
            fontSize=12,
            leading=14,
            spaceAfter=6,
            textColor=colors.darkblue
        ))
        styles.add(ParagraphStyle(
            name="ImageCaption",
            fontSize=10,
            leading=12,
            spaceBefore=6,
            spaceAfter=12,
            textColor=colors.grey,
            fontName=font_name
        ))

        styles.add(ParagraphStyle(
            name="Heading_CN",
            fontName=font_name,
            fontSize=14,
            leading=16,
            spaceAfter=6,
            textColor=colors.darkblue
        ))

        return styles

    def optimize_text(self, text: str) -> str:
        """使用大模型优化文本内容"""
        prompt = ChatPromptTemplate.from_template("""
        你是一位专业的内容优化专家。请优化以下文本，使其更加:
        - 专业且易于理解
        - 结构清晰
        - 语法正确
        - 保持原意但表达更优美

        原始文本:
        {text}

        请直接返回优化后的文本，不要添加任何解释或备注。
        """)

        chain = prompt | self.text_llm | self.output_parser
        return chain.invoke({"text": text})

    @staticmethod
    def analyze_image(image_path: str) -> str:
        """使用多模态模型分析图片并生成描述"""
        return image_analyse_to_text(image_path)

    def create_pdf(self, content: str, image_paths: list = None, output_file: str = "output.pdf"):
        """创建PDF文档"""
        # 1. 设置中文字体

        from reportlab.lib.pagesizes import A4
        from reportlab.lib.units import cm
        # 2. 创建PDF文档
        doc = SimpleDocTemplate(
            output_file,
            pagesize=A4,
            leftMargin=2 * cm,
            rightMargin=2 * cm
        )
        story = []

        # 添加优化后的文本内容
        optimized_content = self.optimize_text(content)
        story.append(Paragraph("优化后的内容", self.styles["Heading_CN"]))
        story.append(Spacer(1, 0.2 * inch))

        for para in optimized_content.split("\n"):
            if para.strip():  # 跳过空行
                story.append(Paragraph(para.strip(), self.styles["Content"]))
                story.append(Spacer(1, 0.1 * inch))

        story.append(PageBreak())

        # 添加图片和图片描述
        if image_paths:
            story.append(Paragraph("图片分析", self.styles["Heading_CN"]))
            story.append(Spacer(1, 0.3 * inch))

            for img_path in image_paths:
                if os.path.exists(img_path):
                    try:
                        # 分析图片
                        img_description = self.analyze_image(img_path)

                        # 添加图片
                        pil_img = PILImage.open(img_path)
                        img_width, img_height = pil_img.size
                        ratio = min(6 * inch / img_width, 8 * inch / img_height)

                        story.append(Image(img_path, width=img_width * ratio, height=img_height * ratio))
                        story.append(Spacer(1, 0.1 * inch))

                        # 添加图片描述
                        story.append(Paragraph(img_description, self.styles["ImageCaption"]))
                        story.append(PageBreak())
                    except Exception as e:
                        print(f"处理图片 {img_path} 时出错: {str(e)}")
                else:
                    print(f"图片不存在: {img_path}")

        # 生成PDF
        doc.build(story)
        print(f"PDF已成功生成: {output_file}")


def setup_chinese_fonts():
    # 获取系统字体目录（通常为 C:\Windows\Fonts）
    fonts_dir = os.path.join(os.environ['WINDIR'], 'Fonts')
    # 微软雅黑
    pdfmetrics.registerFont(TTFont('MSYH', os.path.join(fonts_dir, 'msyh.ttc')))
    return 'MSYH'


# 示例使用
if __name__ == "__main__":
    # 初始化优化器
    optimizer = ContentOptimizer()

    # 示例文本内容
    sample_text = """
    LangChain是一个框架，用于开发由语言模型驱动的应用程序。

    它有很多功能：
    - 可以和不同的语言模型交流
    - 能处理文档
    - 可以把多个功能连起来用
    - 能让模型记住东西
    - 模型可以自己做决定

    用LangChain能做很多东西，比如问答系统、总结文档、生成内容等。
    """

    image_paths_ = [
        './image/lc1.png',
        './image/lc2.png',
        './image/lc3.png',
    ]

    # 过滤掉不存在的图片
    valid_images = [img for img in image_paths_ if os.path.exists(img)]

    # 生成PDF
    optimizer.create_pdf(
        content=sample_text,
        image_paths=valid_images,
        output_file="optimized_report.pdf"
    )
